# 导入所需的库
import matplotlib.pyplot as plt  # 用于绘图
import os  # 用于操作文件和目录
import re  # 用于正则表达式处理
import shutil  # 用于文件和目录的高级操作
import string  # 用于字符串操作
import tensorflow as tf  # 导入 TensorFlow 库

from keras_preprocessing.sequence import pad_sequences

from tensorflow.keras import layers  # 导入 Keras 层模块
from tensorflow.keras import losses  # 导入 Keras 损失函数模块
from tokenizer import tokenizer  # 导入自定义的tokenizer模块（假设该模块在路径下）

print(tf.__version__)  # 输出 TensorFlow 的版本号

# 数据集下载地址
url = "https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz"

# 使用 tf.keras.utils.get_file 下载数据集，并存储到指定路径
dataset = tf.keras.utils.get_file("aclImdb_v1", url,
                                  untar=False, cache_dir=r"E:\tensorflow\tensorflow_learn\datasets",
                                  cache_subdir='', force_download=False)

print(f"download success->{dataset}")  # 输出下载后的数据集路径

# 获取数据集的目录路径
if not os.path.isdir(dataset):
    dataset_dir = os.path.join(os.path.dirname(dataset), 'aclImdb')
else:
    dataset_dir = os.path.join(dataset, 'aclImdb')
print(f"dataset_dir:{dataset_dir}")  # 输出数据集的目录路径

# 获取训练集目录路径
train_dir = os.path.join(dataset_dir, 'train')
# 查看训练集目录下的文件和子目录
print(f"train_dir:{train_dir}")
os.listdir(train_dir)  # 输出训练集文件列表

# 获取测试集目录路径
test_dir = os.path.join(dataset_dir, 'test')
# 查看测试集目录下的文件和子目录
print(f"test_dir:{test_dir}")
os.listdir(test_dir)  # 输出测试集文件列表

# 删除训练集中的 'unsup' 子目录（未标注的数据）
remove_dir = os.path.join(train_dir, 'unsup')
if os.path.exists(remove_dir):
    shutil.rmtree(remove_dir)  # 删除未标注的子目录

# 使用 Keras 提供的 text_dataset_from_directory 函数加载训练集数据
raw_train_ds = tf.keras.utils.text_dataset_from_directory(
    train_dir,
    batch_size=32,
    validation_split=0.2,
    subset='training',
    seed=42)

# 加载验证集数据
raw_val_ds = tf.keras.utils.text_dataset_from_directory(
    train_dir,
    batch_size=32,
    validation_split=0.2,
    subset='validation',
    seed=42)

# 加载测试集数据
raw_test_ds = tf.keras.utils.text_dataset_from_directory(
    test_dir,
    batch_size=32)

# 查看训练集的一个批次的文本和标签
for text_batch, label_batch in raw_train_ds:
    print(text_batch.shape)
    print(label_batch.shape)
    break  # 只打印一个批次的数据

# 准备用于训练的数据集
def custom_standardization(input_data):
    """
    定义一个文本标准化函数：
    1. 将文本转换为小写
    2. 替换 <br /> 标签为空格
    3. 删除标点符号
    """
    lowercase = tf.strings.lower(input_data)  # 转为小写
    stripped_html = tf.strings.regex_replace(lowercase, '<br />', ' ')  # 删除 HTML 标签
    return tf.strings.regex_replace(stripped_html,
                                    '[%s]' % re.escape(string.punctuation),
                                    '')  # 删除标点符号

max_features = 10000  # 最大特征数
sequence_length = 250  # 每个输入序列的最大长度

# 定义 TextVectorization 层，用于将文本转换为整数序列
vectorize_layer = layers.TextVectorization(
    standardize=custom_standardization,  # 使用自定义的标准化函数
    max_tokens=max_features,  # 限制最多的词汇数量
    output_mode='int',  # 输出模式为整数（即词汇表的索引）
    output_sequence_length=sequence_length)  # 设置输出序列的最大长度

# 在训练数据上适应（拟合）vectorize_layer，使其学习文本中的词汇
train_text = raw_train_ds.map(lambda x, y: x)
vectorize_layer.adapt(train_text)

# 定义一个函数，用于将文本和标签进行向量化
def vectorize_text(text, label):
    text = tf.expand_dims(text, -1)  # 增加一个维度，符合 TextVectorization 输入要求
    return vectorize_layer(text), label  # 返回向量化的文本和标签

# 获取一个批次的数据
text_batch, label_batch = next(iter(raw_train_ds))
first_review, first_label = text_batch[0], label_batch[0]
print("Review", first_review)  # 打印第一条评论
print("Label", raw_train_ds.class_names[first_label])  # 打印第一条评论的标签（正面或负面）
print("Vectorized review", vectorize_text(first_review, first_label))  # 打印向量化后的评论

# 查看词汇表中的某些词汇的索引
print("1287 ---> ", vectorize_layer.get_vocabulary()[1287])
print(" 313 ---> ", vectorize_layer.get_vocabulary()[313])
print('Vocabulary size: {}'.format(len(vectorize_layer.get_vocabulary())))  # 输出词汇表大小

# 对训练集、验证集和测试集应用 TextVectorization 层
train_ds = raw_train_ds.map(vectorize_text)
val_ds = raw_val_ds.map(vectorize_text)
test_ds = raw_test_ds.map(vectorize_text)

# 配置数据集，以提高性能
AUTOTUNE = tf.data.AUTOTUNE  # 自动调整数据加载的预取方式
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)  # 缓存并预取数据
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
test_ds = test_ds.cache().prefetch(buffer_size=AUTOTUNE)

# 创建一个简单的模型
embedding_dim = 16  # 嵌入层的维度
model = tf.keras.Sequential([
    layers.Embedding(max_features + 1, embedding_dim),  # 嵌入层，将单词映射为固定长度的向量
    layers.Dropout(0.2),  # Dropout 层，用于减少过拟合
    layers.GlobalAveragePooling1D(),  # 全局平均池化层，将每个特征维度的值取平均
    layers.Dropout(0.2),  # 再次使用 Dropout 层
    layers.Dense(1)  # 输出层，只有一个神经元，表示二分类的结果（正面或负面）
])

model.summary()  # 打印模型概况

# 损失函数与优化器
model.compile(loss=losses.BinaryCrossentropy(from_logits=True),  # 使用二元交叉熵损失函数
              optimizer='adam',  # 使用 Adam 优化器
              metrics=[tf.metrics.BinaryAccuracy(threshold=0.0)])  # 计算二分类的准确度

# 训练模型
epochs = 10  # 训练轮数
history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=epochs)  # 训练模型

# 评估模型
loss, accuracy = model.evaluate(test_ds)  # 在测试集上评估模型

print("Loss: ", loss)
print("Accuracy: ", accuracy)

# 创建准确率和损失随时间变化的图表
history_dict = history.history  # 获取训练历史
history_dict.keys()  # 查看历史记录的键（包含：loss、accuracy 等）

# 获取训练过程中的准确率、验证准确率、训练损失和验证损失
acc = history_dict['binary_accuracy']
val_acc = history_dict['val_binary_accuracy']
loss = history_dict['loss']
val_loss = history_dict['val_loss']

epochs = range(1, len(acc) + 1)  # 获取训练的 epoch 范围

# 绘制训练和验证的损失与准确率图表
plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)
plt.plot(epochs, loss, 'bo', label='Training loss')  # 绘制训练损失
plt.plot(epochs, val_loss, 'b', label='Validation loss')  # 绘制验证损失
plt.title('Training and validation loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(epochs, acc, 'bo', label='Training acc')  # 绘制训练准确率
plt.plot(epochs, val_acc, 'b', label='Validation acc')  # 绘制验证准确率
plt.title('Training and validation accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend(loc='lower right')

plt.show()  # 显示图表

# 导出模型
export_model = tf.keras.Sequential([
  vectorize_layer,  # 文本向量化层，将文本转换为整数序列
  model,  # 之前定义的模型
  layers.Activation('sigmoid')  # 添加 Sigmoid 激活层，将输出映射到 [0, 1] 范围内，适用于二分类任务
])

# 编译模型
export_model.compile(
    loss=losses.BinaryCrossentropy(from_logits=False),  # 使用二元交叉熵损失函数，from_logits=False 表示模型输出已经过 Sigmoid 激活
    optimizer="adam",  # 使用 Adam 优化器
    metrics=['accuracy']  # 评估指标为准确率
)

# 使用原始测试集进行评估
loss, accuracy = export_model.evaluate(raw_test_ds)
print(accuracy)  # 输出模型在测试集上的准确率

# 使用新数据进行推断
examples = [
  "The movie was great!",  # 示例评论 1
  "The movie was okay.",   # 示例评论 2
  "The movie was terrible..."  # 示例评论 3
]

# 使用模型进行预测
predictions = export_model.predict(examples)
print(predictions)  # 输出模型对示例评论的预测结果
